#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据清洗和优化工具
功能：清洗决策树数据，提高向量化搜索命中率
"""

import re
import jieba
import pandas as pd
from typing import List, Dict, Tuple


class DataCleaner:
    def __init__(self):
        # 常见的同义词映射
        self.synonyms = {
            '无产物': ['没有产物', '无条带', '扩增失败', '无扩增产物'],
            '产物量少': ['产物量很少', '条带弱', '扩增效率低', '产出低'],
            '模板质量': ['模板品质', 'DNA质量', 'RNA质量', '模板完整性'],
            '投入量': ['用量', '加入量', '模板量', '起始量'],
            '降解': ['破碎', '断裂', '不完整', '质量差'],
            'PCR反应': ['扩增反应', 'PCR扩增', '聚合酶反应'],
            '建议': ['推荐', '应该', '可以', '需要']
        }
        
        # 技术术语标准化
        self.term_standardization = {
            'μl': 'ul',
            'ng': 'ng',
            'bp': 'bp',
            'DNA': 'DNA',
            'RNA': 'RNA',
            'cDNA': 'cDNA',
            'gDNA': 'gDNA',
            'PCR': 'PCR'
        }
    
    def parse_decision_tree(self, text: str) -> Dict:
        """解析决策树结构"""
        # 分割路径
        parts = text.split(' -> ')
        
        if len(parts) < 2:
            return None
        
        # 提取产品编号
        product = parts[0].strip()
        
        # 提取货号
        cargo_match = re.search(r'\[货号:([^\]]+)\]', text)
        cargo_number = cargo_match.group(1) if cargo_match else product
        
        # 构建路径
        path = ' -> '.join(parts[1:-1]) if len(parts) > 2 else ''
        
        # 提取最终结果/建议
        result = parts[-1]
        if cargo_match:
            result = result.replace(cargo_match.group(0), '').strip()
        
        return {
            'product': product,
            'cargo_number': cargo_number,
            'problem_path': path,
            'solution': result,
            'full_text': text
        }
    
    def extract_keywords(self, text: str) -> List[str]:
        """提取关键词"""
        # 分词
        words = list(jieba.cut(text))
        
        # 过滤和清理
        keywords = []
        for word in words:
            word = word.strip()
            if len(word) > 1 and not word.isdigit():
                # 标准化术语
                word = self.term_standardization.get(word, word)
                keywords.append(word)
        
        return keywords
    
    def generate_search_variants(self, parsed_data: Dict) -> List[str]:
        """生成搜索变体，提高命中率"""
        variants = []
        
        # 1. 原始完整文本
        variants.append(parsed_data['full_text'])
        
        # 2. 问题描述 + 解决方案
        if parsed_data['problem_path'] and parsed_data['solution']:
            variants.append(f"{parsed_data['problem_path']} -> {parsed_data['solution']}")
        
        # 3. 产品 + 问题
        if parsed_data['problem_path']:
            variants.append(f"{parsed_data['product']} {parsed_data['problem_path']}")
        
        # 4. 产品 + 解决方案
        variants.append(f"{parsed_data['product']} {parsed_data['solution']}")
        
        # 5. 关键词组合
        keywords = self.extract_keywords(parsed_data['full_text'])
        if len(keywords) >= 3:
            # 取前3个关键词
            key_combo = ' '.join(keywords[:3])
            variants.append(key_combo)
        
        # 6. 添加同义词变体
        for original, synonyms in self.synonyms.items():
            if original in parsed_data['full_text']:
                for synonym in synonyms:
                    variant_text = parsed_data['full_text'].replace(original, synonym)
                    variants.append(variant_text)
        
        return variants
    
    def extract_structured_info(self, parsed_data: Dict) -> Dict:
        """提取结构化信息"""
        info = {
            'product_code': parsed_data['product'],
            'cargo_number': parsed_data['cargo_number'],
            'problem_category': '',
            'problem_detail': '',
            'condition': '',
            'solution': parsed_data['solution'],
            'parameters': {}
        }
        
        # 解析问题路径
        path_parts = parsed_data['problem_path'].split(' -> ')
        if len(path_parts) >= 1:
            info['problem_category'] = path_parts[0]
        if len(path_parts) >= 2:
            info['problem_detail'] = path_parts[1]
        if len(path_parts) >= 3:
            info['condition'] = ' -> '.join(path_parts[2:])
        
        # 提取数值参数
        numbers = re.findall(r'(\d+(?:\.\d+)?)\s*([a-zA-Zμ]+)', parsed_data['full_text'])
        for value, unit in numbers:
            info['parameters'][f'{value}{unit}'] = f'{value} {unit}'
        
        return info
    
    def clean_and_enhance_data(self, raw_lines: List[str]) -> List[Dict]:
        """清洗和增强数据"""
        enhanced_data = []
        
        for line in raw_lines:
            line = line.strip()
            if not line:
                continue
            
            # 解析决策树结构
            parsed = self.parse_decision_tree(line)
            if not parsed:
                continue
            
            # 生成搜索变体
            search_variants = self.generate_search_variants(parsed)
            
            # 提取结构化信息
            structured_info = self.extract_structured_info(parsed)
            
            # 为每个变体创建记录
            for i, variant in enumerate(search_variants):
                enhanced_record = {
                    'id': f"{parsed['product']}_{hash(line)}_{i}",
                    'original_text': line,
                    'search_text': variant,
                    'product_code': structured_info['product_code'],
                    'cargo_number': structured_info['cargo_number'],
                    'problem_category': structured_info['problem_category'],
                    'problem_detail': structured_info['problem_detail'],
                    'condition': structured_info['condition'],
                    'solution': structured_info['solution'],
                    'parameters': str(structured_info['parameters']),
                    'variant_type': ['original', 'problem_solution', 'product_problem', 
                                   'product_solution', 'keywords', 'synonym'][i] if i < 6 else 'synonym',
                    'keywords': ' '.join(self.extract_keywords(variant))
                }
                enhanced_data.append(enhanced_record)
        
        return enhanced_data
    
    def export_to_csv(self, enhanced_data: List[Dict], output_file: str):
        """导出到CSV文件"""
        df = pd.DataFrame(enhanced_data)
        df.to_csv(output_file, index=False, encoding='utf-8')
        print(f"✅ 增强数据已导出到: {output_file}")
        print(f"📊 原始记录: {len(set(d['original_text'] for d in enhanced_data))}")
        print(f"📊 增强记录: {len(enhanced_data)}")
        
        # 统计信息
        variant_counts = df['variant_type'].value_counts()
        print("\n📈 变体类型分布:")
        for variant_type, count in variant_counts.items():
            print(f"  {variant_type}: {count}")


def demo_cleaning():
    """演示数据清洗"""
    # 示例数据
    sample_data = [
        "PK511 -> 无产物或产物量很少 -> 模板 -> 是否鉴定过模板质量 -> 是 -> 质量较好，则没有问题 [货号:PK511]",
        "PK511 -> 无产物或产物量很少 -> 模板 -> 是否鉴定过模板质量 -> 是 -> 模板降解，则可能扩增不出条带 -> 建议重新制备模板 [货号:PK511]",
        "PK511 -> 无产物或产物量很少 -> 模板 -> 是否鉴定过模板质量 -> 否 -> 建议鉴定模板质量 -> cDNA不可直接跑胶，可以跑胶鉴定RNA的质量 [货号:PK511]",
        "PK511 -> 无产物或产物量很少 -> 模板 -> 50 μl投入量 -> 基因组DNA -> ＞500ng，则可能抑制PCR反应 -> 建议降低投入量 [货号:PK511]",
        "PK511 -> 无产物或产物量很少 -> 模板 -> 50 μl投入量 -> 基因组DNA -> 5-500 ng -> 没有问题 [货号:PK511]"
    ]
    
    # 创建清洗器
    cleaner = DataCleaner()
    
    # 清洗和增强数据
    enhanced_data = cleaner.clean_and_enhance_data(sample_data)
    
    # 导出结果
    cleaner.export_to_csv(enhanced_data, 'enhanced_pk511_data.csv')
    
    # 显示一些示例
    print("\n📝 增强数据示例:")
    for i, record in enumerate(enhanced_data[:5]):
        print(f"\n--- 记录 {i+1} ---")
        print(f"变体类型: {record['variant_type']}")
        print(f"搜索文本: {record['search_text'][:100]}...")
        print(f"问题分类: {record['problem_category']}")
        print(f"解决方案: {record['solution'][:50]}...")


if __name__ == "__main__":
    demo_cleaning()
